#include <vector>
#include <iomanip> 
#include <iostream>
#include <chrono>
#include "matrix_operations_gpu.cu"
#include "matrix_operations_cpu.h"
#include "matrix_operations_wmma.cu"
#include "simple_fnn.h"
#include "relu.h"
#include "sigmoid.h"
#include "loss_function.cu"
#include <time.h>

bool areMatricesEqual(const std::vector<std::vector<float>>& A, const std::vector<std::vector<float>>& B, float tolerance = 1e-6) {
    if (A.size() != B.size() || A[0].size() != B[0].size()) {
        return false;
    }

    for (size_t i = 0; i < A.size(); ++i) {
        for (size_t j = 0; j < A[i].size(); ++j) {
            if (std::abs(A[i][j] - B[i][j]) > tolerance) {
                return false;
            }
        }
    }
    return true;
}

void printMatrix(const std::vector<std::vector<float>>& matrix) {
    for (const auto& row : matrix) {
        for (float elem : row) {
            std::cout << elem << " ";
        }
        std::cout << std::endl;
    }
}
void compareMatrix(const std::vector<std::vector<float>>& matrix_A,const std::vector<std::vector<float>>& matrix_B){
    if (matrix_A.size() != matrix_B.size() || matrix_A[0].size() != matrix_B[0].size()) {
        printf("don't match");
    }
    int error=0;
    for (size_t i = 0; i < matrix_A.size(); ++i) {
        for (size_t j = 0; j < matrix_A[i].size(); ++j) {
            if (std::abs(matrix_A[i][j] - matrix_B[i][j]) > 0.001) {
                //std::cout<<"error at row "<<i<<" col "<<j<<std::endl;
                error++;
            }
        }
    }
    std::cout<<"total:"<<(matrix_A.size()*matrix_A.size())<<" error:"<<error<<" rate:"<<float(error)/(matrix_A.size()*matrix_A.size())<<std::endl;
}
int main(){
     srand(time(NULL));
    // Example matrices
    const int size = 16;  // Set the size of the matrix to 16x16
    int M = size*1024;
    int N = size*1024;
    int K = size*2;
    std::vector<std::vector<float>> A(M, std::vector<float>(K));
    std::vector<std::vector<float>> B(K, std::vector<float>(N));

    // Fill matrices with some values
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < K; ++j) {
            A[i][j] = 1;  

        }
    }
    for (int i = 0; i < K; ++i) {
        for (int j = 0; j < N; ++j) {
            B[i][j] = 1;  
        }
    }


    MatrixOperationsCPU cpu;
    MatrixOperationsGPU gpu;
    MatrixOperationsWMMA wmma;

    // Performance testing
    // auto startCPU = std::chrono::high_resolution_clock::now();
    // auto resultCPU = cpu.multiply(A, B);
    // auto endCPU = std::chrono::high_resolution_clock::now();

    // auto startGPU = std::chrono::high_resolution_clock::now();
    // auto resultGPU = gpu.multiply(A, B);
    // auto endGPU = std::chrono::high_resolution_clock::now();

    auto startWMMA = std::chrono::high_resolution_clock::now();
    auto resultWMMA = wmma.multiply(A, B);
    auto endWMMA = std::chrono::high_resolution_clock::now();

    // std::chrono::duration<double, std::milli> cpuTime = endCPU - startCPU;
    // std::chrono::duration<double, std::milli> gpuTime = endGPU - startGPU;
    std::chrono::duration<double, std::milli> wmmaTime = endWMMA - startWMMA;

    // std::cout << "CPU Result:" << std::endl;
    // printMatrix(resultCPU);
    // std::cout << "GPU Result:" << std::endl;
    // printMatrix(resultGPU);

    // std::cout << "CPU Time: " << cpuTime.count() << " ms" << std::endl;
    // std::cout << "GPU Time: " << gpuTime.count() << " ms" << std::endl;
    std::cout << "WMMA Time: " << wmmaTime.count() << " ms" << std::endl;
    // std::cout << "Matrix A:" << std::endl;
    // printMatrix(A);

    // std::cout << "Matrix B:" << std::endl;
    // printMatrix(B);

    // Compute on CPU and GPU
    // auto resultCPU = cpu.multiply(A, B);
    // auto resultGPU = gpu.multiply(A, B); 
    // auto resultWMMA = wmma.multiply(A, B); 

    // std::cout << "CPU Result:" << std::endl;
    // printMatrix(resultCPU);


    // printMatrix(resultGPU);

    //std::cout << "WMMA Result:" << std::endl;
    //printMatrix(resultWMMA);
    
    //compareMatrix(resultGPU,resultWMMA);

}